import os
import time
from datetime import datetime
from typing import Optional
import numpy as np
from einops import rearrange
import json
from .data_config import DataManager, DataConfig


def rseq(file_path: str) -> dict:
    """
    Load a JSON file and return its content.

    Parameters:
    - file_path (str): The path to the JSON file.

    Returns:
    - dict: The content of the JSON file.
    """
    if not file_path.endswith('.json'):
        file_path += '.json'

    with open(file_path, 'r') as f:
        content = json.load(f)

    return content


class ModalityDataGuide:
    def __init__(self, modality: DataConfig, mask_guide=True):
        root = modality.data_path
        self.img_size = modality.img_size
        self.patch_size = modality.patch_size
        self.grid_size = modality.grid_size
        self.token_dim = modality.token_dim
        self.num_channels = modality.num_channels
        self.mask_guide = mask_guide
        if root == '':
            self.data_list = [None for _ in range(50)]  # for testing
        else:
            self.data_list = []
            for i in [os.path.join(root, p) for p in os.listdir(root) if '.json' in p]:
                self.data_list += rseq(i).get('files')
            # self.data_list = [os.path.join(root, p) for p in os.listdir(root)]
        print(f'Modality {modality.name}: {len(self.data_list)} instances')

    def patchify(self, d):
        p, g, c = self.patch_size, self.grid_size, self.num_channels
        if len(p) == 3:
            d = rearrange(
                d,
                'c (g0 p0) (g1 p1) (g2 p2) -> (g0 g1 g2) (c p0 p1 p2)',
                c=c, g0=g[0], g1=g[1], g2=g[2], p0=p[0], p1=p[1], p2=p[2]
            )
        elif len(p) == 2:
            d = rearrange(
                d,
                'c (g0 p0) (g1 p1) -> (g0 g1) (c p0 p1)',
                c=c, g0=g[0], g1=g[1], p0=p[0], p1=p[1],
            )
        elif len(p) == 1:
            d = rearrange(
                d,
                'c (g0 p0) -> g0 (c p0)',
                c=c, g0=g[0], p0=p[0],
            )
        return d

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, idx):
        data = self.data_list[idx]
        if isinstance(data, str):
            img = np.load(os.path.join(data, 'img.npy')).astype(np.float16)
            if self.num_channels == 1:
                img = img[None, :]

            msk_path = os.path.join(data, 'msk.npy')
            if os.path.exists(msk_path) and self.mask_guide:
                msk = np.load(msk_path)[None, :].astype(np.float16)
            else:
                msk = np.ones((1, *self.img_size), dtype=np.float16)

            # ---- adapt to combined img and msk
            # if self.mask_guide:
            #     cim = np.load(os.path.join(data, 'cim.npy')).astype(np.float16)
            #     img = cim[0, :]
            #     msk = cim[1, :]
            #     if self.num_channels == 1:
            #         img = img[None, :].astype(np.float16)
            #         msk = msk[None, :].astype(np.float16)
            # else:
            #     img = np.load(os.path.join(data, 'img.npy')).astype(np.float16)
            #     if self.num_channels == 1:
            #         img = img[None, :]
            #     msk = np.ones((1, *self.img_size), dtype=np.float16)
        else:
            img = np.random.rand(self.num_channels, *self.img_size).astype(np.float16)
            msk = np.random.randint(0, 2, size=(1, *self.img_size))
        # make msk per pixel

        msk = np.broadcast_to(msk, (self.num_channels, *self.img_size))
        img = self.patchify(img)
        msk = self.patchify(msk)
        return img, msk


class MultiModalDatasetGuide:
    column_names = [
        'image', 'mask_guide', 'modality_idx',
    ]

    def __init__(self, 
                 data_manager: DataManager, 
                 single_modality: Optional[str] = None, 
                 balance: bool = False, 
                 fix_batches: int = -1, 
                 mask_guide: bool = True):
        if mask_guide:
            print('[Dataset] Use Mask Guide!!!!')

        self.data_manager = data_manager
        self.mask_guide = mask_guide
        self.n_modality = self.data_manager.n_modality
        if single_modality is not None:
            self.datasets = []
            for m in data_manager:
                if m.name == single_modality:
                    self.datasets.append(ModalityDataGuide(m, mask_guide))
                else:
                    self.datasets.append([])
        else:
            self.datasets = [ModalityDataGuide(m, mask_guide) for m in data_manager]
        self.datasets_len = [len(d) for d in self.datasets]

        self.balance = balance
        self.fix_batches = fix_batches

        max_datasets_len = max(self.datasets_len)
        if balance or self.fix_batches > 0:
            if self.fix_batches > 0 and self.fix_batches < max_datasets_len:
                print(
                    f'Warning: Fix batches ({self.fix_batches}) is smaller than max modality batches ({max_datasets_len})!')
            target_len = fix_batches if fix_batches > 0 else max_datasets_len
            data_list_c1 = np.concatenate([
                np.ones(target_len) * d_idx for d_idx, d in enumerate(self.datasets)
            ])
            data_list_c2 = np.concatenate([
                np.arange(target_len) % len(d) for d in self.datasets
            ])
        else:
            data_list_c1 = np.concatenate([np.ones(len(d)) * d_idx for d_idx, d in enumerate(self.datasets)])
            data_list_c2 = np.concatenate([np.arange(len(d)) for d in self.datasets])

        self.data_list = np.concatenate([data_list_c1.reshape(-1, 1), data_list_c2.reshape(-1, 1)], axis=1).astype(int)

        self.total_dim = data_manager.total_dim
        self.seq_len = data_manager.seq_len

    def __len__(self):
        return self.data_list.shape[0]

    def __getitem__(self, idx):
        m_idx, data_idx = self.data_list[idx]
        m = self.data_manager[m_idx]
        m_seq_len = m.seq_len
        lo, hi = m.get_range()
        ## load data
        data_img = np.zeros((self.seq_len, self.total_dim), dtype=np.float16)
        data_msk = np.zeros((self.seq_len, self.total_dim), dtype=np.float16)

        img, msk = self.datasets[m_idx][data_idx]
        data_img[:m_seq_len, lo:hi] = img
        data_msk[:m_seq_len, lo:hi] = msk
        return data_img, data_msk, m_idx


class MultiModalMAEMaskGuide:
    column_names = [
        'image', 'modality_idx', 'instance_ids',
        'target_mask', 'rand_indices', 'unmask_idx', 'mask_guide'
    ]

    def __init__(self, 
                 data_manager: DataManager, 
                 mask_ratio: float = 0.75, 
                 seed: Optional[int] = None, 
                 mask_guide: bool =True, 
                 guide_mode: str = 'patch',  
                 guide_scale: float = 2.0, 
                 num_prefix_tokens: int = 1):
        if mask_guide:
            print('[Mask] Use Mask Guide!!!!')

        self.mask_guide = mask_guide
        self.guide_mode = guide_mode
        self.guide_scale = guide_scale
        self.data_manager = data_manager
        self.n_modality = self.data_manager.n_modality

        self.total_dim = data_manager.total_dim
        self.seq_len = data_manager.seq_len
        self.pad_lens = data_manager.pad_lens
        self.m_seq_lens = data_manager.m_seq_lens

        self.keep_num = int((1 - mask_ratio) * self.seq_len)
        # precompute instance ids
        # plus one for cls token
        self.m_instance_ids = np.zeros((self.n_modality, num_prefix_tokens + self.seq_len), dtype=np.int)
        for m_idx in range(self.n_modality):
            self.m_instance_ids[m_idx, :self.m_seq_lens[m_idx] + num_prefix_tokens] = 1
        self.seed = seed

    def __call__(self, data, msk, m_idx):
        m = self.data_manager[m_idx]
        m_seq_len = m.seq_len
        lo, hi = m.get_range()
        pad_len = self.pad_lens[m_idx]

        ## add random mask
        if self.seed is not None:
            state = np.random.RandomState(seed=self.seed)
            noise = state.uniform(size=(self.seq_len,))
        else:
            noise = np.random.uniform(size=(self.seq_len,))

        rand_indices = np.argsort(noise, axis=0).astype(np.int32)

        unmask_index = rand_indices[ :self.keep_num]
        target_token_index = rand_indices[self.keep_num:]

        target_mask = np.zeros((self.seq_len, self.total_dim), dtype=np.float16)

        # instance id, 0 for pad and invalid tokens
        instance_ids = self.m_instance_ids[m_idx]

        target_mask[target_token_index, lo:hi] = 1
        if pad_len > 0:
            target_mask[-pad_len:, :] = 0

        return data, m_idx, instance_ids, target_mask, rand_indices, unmask_index, msk
